import pybuda._C.graph
from typing import List, Union, overload

class AutogradConfig:
    def __init__(self, recompute: bool = ..., optimizer: object = ...) -> None: ...

class AutogradContext:
    def __init__(self, *args, **kwargs) -> None: ...
    @overload
    def constant(self, arg0: int) -> pybuda._C.graph.NodeContext: ...
    @overload
    def constant(self, arg0: float) -> pybuda._C.graph.NodeContext: ...
    def create_optimizer_op(self, type: str, operands: List[pybuda._C.graph.NodeContext], attributes=...) -> pybuda._C.graph.NodeContext: ...
    def get_operands(self, arg0: pybuda._C.graph.NodeContext) -> List[pybuda._C.graph.NodeContext]: ...
    def get_pytorch_tensor(self, arg0: pybuda._C.graph.NodeContext) -> object: ...
    def get_shape(self, arg0: pybuda._C.graph.NodeContext) -> List[int]: ...
    def input(self, *args, **kwargs): ...
    def loopback(self, arg0: pybuda._C.graph.NodeContext, arg1: pybuda._C.graph.NodeContext) -> None: ...
    def op(self, type: Union[str, object], operands: List[pybuda._C.graph.NodeContext], attributes=...) -> pybuda._C.graph.NodeContext: ...
    def tensor(self, arg0: object) -> pybuda._C.graph.NodeContext: ...

class AutogradEngine:
    def __init__(self, arg0: pybuda._C.graph.Graph, arg1: AutogradConfig) -> None: ...
    def run(self) -> pybuda._C.graph.Graph: ...
